import os
import socket
import torch
import torch.distributed as dist


def init_distributed_mode(conf, rank):

    # 创建进程组  (windows不支持NCCL backend)
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)  # 无线局域网适配器WLAN所对应的IPV4地址
    s.connect(("8.8.8.8", 80))
    os.environ['MASTER_ADDR'] = s.getsockname()[0]
    os.environ['MASTER_PORT'] = '29500'
    # 初始化方法为默认env
    dist.init_process_group(backend='gloo',  # 通信后端默认使用gloo
                            world_size=conf['world_size'],
                            rank=rank)  # 不同进程rank不同
    dist.barrier()  # 等待每一个进程都运行到此再往下运行


def cleanup():  # 关闭进程组并去初始化
    dist.destroy_process_group()


def is_dist_avail_and_initialized():
    """检查是否支持分布式环境"""
    if not dist.is_available():  # 如果分布式包可用，则返回True
        return False
    if not dist.is_initialized():  # 检查默认进程组是否已经初始化
        return False
    return True


def get_world_size():
    if not is_dist_avail_and_initialized():
        return 1
    return dist.get_world_size()  # 返回当前进程组的进程数目。


def get_rank():
    if not is_dist_avail_and_initialized():
        return 0
    return dist.get_rank()  # 返回所提供进程组的当前进程的rank


def is_main_process():  # 通过get_rank() == 0是否成立，判断是否为主进程
    return get_rank() == 0


def reduce_value(value, average=True):
    world_size = get_world_size()  # 得到world_size的大小
    if world_size < 2:  # 单进程(GPU)的情况，则直接返回value，即该函数对于单进程不起作用
        return value

    with torch.no_grad():
        dist.all_reduce(value)  # 对不同设备(GPU)的value进行求和操作
        # AllReduce是一种将所有process中的目标数组（即表示All），减少(运算)为单个数组（即表示Reduce）并将结果数组返回给所有process的操作。
        # (比如将所有GPU上的梯度值，假设数组表示，合并并执行reduce操作成一个数组，并返回给所有GPU)
        if average:  # 求平均
            value /= world_size

        return value
